import java.util.TreeMap;

public class LeetCode347 {
    private class Freq implements Comparable<Freq>{
        public int e,freq;

        public Freq(int e, int freq) {
            this.e = e;
            this.freq = freq;
        }

        @Override
        public int compareTo(Freq another) {
            if (this.freq < another.freq){
                return 1;
            } else if (this.freq > another.freq) {
                return -1;
            }
            return 0;
        }
    }
    public int[] topKFrequent(int[] nums, int k) {
        TreeMap<Integer, Integer> map = new TreeMap<>();

        for (int num:nums) {
            if (map.containsKey(num)){
                map.put(num,map.get(num)+1);
            }else {
                map.put(num,1);
            }
        }
        PriorityQueue<Freq> pq = new PriorityQueue<>();
        for (int key:map.keySet()){
            if (pq.getSize() < k){
                pq.enqueue(new Freq(key, map.get(key)));
            } else if (map.get(key) > pq.getFront().freq){
                pq.dequeue();
                pq.enqueue(new Freq(key, map.get(key)));
            }
        }
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = pq.dequeue().e;
        }

        return res;
    }

}
